/*
 * Copyright The Slam Authors
 */

#include "slam/mapping/reflector_detector.h"

#include <Eigen/Core>

namespace slam {
namespace mapping {

void ComputeReflectorCoordinateTransform(
    const std::vector<sensor::MatchedReflectorInfo>& matched_reflector_info,
    const transform::Rigid2d& init_pose, transform::Rigid2d* matched_pose) {
  double T_map_base[3] = {init_pose.translation().x(),
                          init_pose.translation().y(),
                          init_pose.rotation().angle()};

  ceres::Problem::Options problem_options;
  ceres::Problem problem(problem_options);
  problem.AddParameterBlock(T_map_base, 3);
  for (const auto& item : matched_reflector_info) {
    problem.AddResidualBlock(
        CoordinateTransformCost2D::CreateAutoDiffCostFunction(
            item.reflector_in_tracking, item.reflector_in_map, 10),
        nullptr, T_map_base);
  }
  problem.AddResidualBlock(
      mapping::optimization::DeltaPoseCostFunctor2D::CreateAutoDiffCostFunction(
          1.0, 1.0, T_map_base),
      nullptr, T_map_base);
  ceres::Solver::Options options;
  options.max_num_iterations = 20;
  options.linear_solver_type = ceres::DENSE_QR;
  options.minimizer_progress_to_stdout = false;
  ceres::Solver::Summary summary;
  ceres::Solve(options, &problem, &summary);
  *matched_pose =
      transform::Rigid2d({T_map_base[0], T_map_base[1]}, T_map_base[2]);
}

ReflectorDetect::ReflectorDetect(const proto::ReflectorDectorOptions& options)
    : options_(options) {
  LineExtractionParams line_params;
  const double min_line_length =
      options_.reflector_length() - options_.length_error();
  line_params.min_line_length = min_line_length > 0 ? min_line_length : 0;
  line_params.max_line_length =
      options_.reflector_length() + options_.length_error();
  line_params.max_line_gap = options_.max_line_gap();
  line_params.min_split_dist = options_.min_line_split_dist();
  line_params.outlier_dist = options_.line_outlier_dist();
  line_params.min_line_points = options_.min_reflector_counter_per_cluster();
  if (options_.reflector_type() == ReflectorType::REFLECTOR_PLANE) {
    line_extraction_ = absl::make_unique<LineExtraction>(line_params);
  }
}

void ReflectorDetect::ExtractPlaneFeature(const sensor::PointCloud& points,
                                          sensor::PointCloud* centers_ptr,
                                          std::vector<int>* count_ptr) {
  sensor::PointCloud& centers = *CHECK_NOTNULL(centers_ptr);
  std::vector<int>& count = *CHECK_NOTNULL(count_ptr);
  sensor::PointCloud center_sums;
  std::vector<sensor::PointCloud> cluster_points;
  const int min_points_num = options_.min_reflector_counter_per_cluster();
  if (points.size() < min_points_num) return;

  std::vector<LineInfo> lines_info;
  line_extraction_->ExtractLines(points, &lines_info);
  for (const auto& item : lines_info) {
    sensor::RangefinderPoint center_point;
    center_point.position = Eigen::Vector3f(item.center[0], item.center[1],
                                            points.front().position[2]);
    center_point.intensity = points.front().intensity;
    centers.push_back(center_point);
    count.push_back(item.num_points);
#if 0
    LOG_EVERY_N(WARNING,100) << "[line info] \n\tline_start: " << item.start[0] << ", "
                 << item.start[1] << ", \n\tline_end: " << item.end[0] << " "
                 << item.end[1] << ", \n\tline_center: " << item.center[0]
                 << " " << item.center[1]
                 << ", \n\tnum_points: " << item.num_points
                 << ", \n\tline_length: " << item.length
                 << ", \n\tline_angle: " << item.angle
                 << ", \n\tline_radius: " << item.radius;
#endif
  }
}

void ReflectorDetect::ExtractColumnFeature(const sensor::PointCloud& points,
                                           sensor::PointCloud* centers_ptr,
                                           std::vector<int>* count_ptr) {
  sensor::PointCloud& centers = *CHECK_NOTNULL(centers_ptr);
  std::vector<int>& count = *CHECK_NOTNULL(count_ptr);
  sensor::PointCloud center_sums;
  std::vector<sensor::PointCloud> cluster_points;

  const int min_points_num = options_.min_reflector_counter_per_cluster();
  if (points.size() < min_points_num) return;

  centers.push_back(points[0]);
  center_sums.push_back(points[0]);
  count.push_back(1);

  cluster_points.resize(1);
  cluster_points.at(0).push_back(points[0]);

  for (int i = 1; i < static_cast<int>(points.size()); i++) {
    const Eigen::Vector3f& pt = points[i].position;
    int k = -1;
    for (int j = 0; j < static_cast<int>(centers.size()); j++) {
      const float dist = (pt - centers[j].position).norm();

      if (dist < options_.max_line_gap()) {
        center_sums[j].position += pt;
        count[j]++;

        centers[j].position = (1.0 / count[j]) * center_sums[j].position;
        k = j;

        if (j >= cluster_points.size()) {
          cluster_points.resize(j + 1);
        }
        cluster_points.at(j).push_back(points[i]);
        break;
      }
    }

    if (k < 0) {
      centers.push_back(points[i]);
      center_sums.push_back(points[i]);
      count.push_back(1);
    }
  }

  if (options_.circle_fit_method() == CircleFitMethod::LEAST_SQUARE) {
    LeastSquareCircle(cluster_points, centers_ptr, count_ptr);
  } else {
    RansacCircle(cluster_points, centers_ptr, count_ptr);
  }
}

void ReflectorDetect::ClusterPoints(const sensor::PointCloud& points,
                                    sensor::PointCloud* centers_ptr,
                                    std::vector<int>* count_ptr) {
  if (options_.reflector_type() == ReflectorType::REFLECTOR_PLANE) {
    ExtractPlaneFeature(points, centers_ptr, count_ptr);
  } else if (options_.reflector_type() == ReflectorType::REFLECTOR_COLUMN) {
    ExtractColumnFeature(points, centers_ptr, count_ptr);
  }
}

void ReflectorDetect::RansacCircle(
    const std::vector<sensor::PointCloud>& cluster_points,
    sensor::PointCloud* centers_ptr, std::vector<int>* count_ptr) {
  sensor::PointCloud& centers = *CHECK_NOTNULL(centers_ptr);
  std::vector<int>& count = *CHECK_NOTNULL(count_ptr);
  const int min_points_num =
      std::max(options_.min_reflector_counter_per_cluster(), 3);

  sensor::PointCloud column_centers;
  std::vector<int> column_points_count;

  for (int i = 0; i < cluster_points.size(); i++) {
    const sensor::PointCloud& cluster_point = cluster_points.at(i);
    const int points_num = cluster_point.size();
    if (points_num < min_points_num) continue;

    Eigen::Vector3f& center_position = centers.at(i).position;

    const int max_iter_cnt = std::max(points_num, 5);
    const double max_error_threshold = options_.radius_error();
    const int consensus_cnt_threshold = min_points_num;

    std::default_random_engine rng;
    std::uniform_int_distribution<int> uniform(0, points_num - 1);
    rng.seed(777);

    std::vector<int> best_consensus_indexs;
    double center_x = 0., center_y = 0., radius = 0.;
    double model_mean_error = 0.;
    int best_consensus_cnt = 0;
    int iter = 0;

    while (iter < max_iter_cnt) {
      std::vector<Eigen::Vector3f> select_points;
      for (int c = 0; c < 3; ++c) {
        select_points.push_back(cluster_point.at(uniform(rng)).position);
      }

      const double& x1 = select_points.at(0)[0];
      const double& y1 = select_points.at(0)[1];

      const double& x2 = select_points.at(1)[0];
      const double& y2 = select_points.at(1)[1];

      const double& x3 = select_points.at(2)[0];
      const double& y3 = select_points.at(2)[1];

      double a = x1 - x2;
      double b = y1 - y2;
      double c = x1 - x3;
      double d = y1 - y3;
      double e = ((x1 * x1 - x2 * x2) + (y1 * y1 - y2 * y2)) / 2.0;
      double f = ((x1 * x1 - x3 * x3) + (y1 * y1 - y3 * y3)) / 2.0;
      double det = b * c - a * d;

      if (fabs(det) > 1e-5) {
        center_x = -(d * e - b * f) / det;
        center_y = -(a * f - c * e) / det;
        radius = (hypot(x1 - center_x, y1 - center_y) +
                  hypot(x2 - center_x, y2 - center_y) +
                  hypot(x3 - center_x, y3 - center_y)) /
                 3.0;

        double mean_error = 0.;
        std::vector<int> tmp_consensus_indexs;
        tmp_consensus_indexs.reserve(points_num);

        for (int j = 0; j < points_num; ++j) {
          const Eigen::Vector3f& point = cluster_point.at(j).position;
          const double distance =
              abs(radius - sqrt((point[0] - center_x) * (point[0] - center_x) +
                                (point[1] - center_y) * (point[1] - center_y)));

          if (distance < max_error_threshold) {
            tmp_consensus_indexs.push_back(j);
          }
          mean_error += distance;
        }

        if (tmp_consensus_indexs.size() >= best_consensus_cnt &&
            tmp_consensus_indexs.size() >= consensus_cnt_threshold) {
          model_mean_error = mean_error / points_num;
          best_consensus_indexs.clear();
          best_consensus_indexs = tmp_consensus_indexs;
          best_consensus_cnt = best_consensus_indexs.size();
          center_position[0] = center_x;
          center_position[1] = center_y;
        }
      }
      iter++;
    }

    if (radius > options_.reflector_radius() - options_.radius_error() &&
        radius < options_.reflector_radius() + options_.radius_error()) {
      column_centers.emplace_back(centers.at(i));
      column_points_count.emplace_back(count.at(i));
    } else {
      LOG_EVERY_N(WARNING, 10) << "radius: " << radius;
    }
  }

  centers.clear();
  count.clear();
  centers = column_centers;
  count = column_points_count;
}

void ReflectorDetect::LeastSquareCircle(
    const std::vector<sensor::PointCloud>& cluster_points,
    sensor::PointCloud* centers_ptr, std::vector<int>* count_ptr) {
  sensor::PointCloud& centers = *CHECK_NOTNULL(centers_ptr);
  std::vector<int>& count = *CHECK_NOTNULL(count_ptr);
  const int min_points_num =
      std::max(options_.min_reflector_counter_per_cluster(), 3);

  sensor::PointCloud column_centers;
  std::vector<int> column_points_count;

  for (int i = 0; i < cluster_points.size(); i++) {
    const sensor::PointCloud& cluster_point = cluster_points.at(i);
    const int points_num = cluster_point.size();
    if (points_num < min_points_num) continue;

    Eigen::Vector3f& center_position = centers.at(i).position;

    double init_para[3] = {center_position[0], center_position[1],
                           std::sqrt(options_.reflector_radius())};
    ceres::Problem::Options problem_options;
    ceres::Problem problem(problem_options);
    problem.AddParameterBlock(init_para, 3);
    ceres::LossFunction* huber_loss =
        new ceres::HuberLoss(options_.radius_error());
    for (int j = 0; j < points_num; ++j) {
      const Eigen::Vector3f& point = cluster_point.at(j).position;
      problem.AddResidualBlock(
          DistanceFromCircleCost::CreateAutoDiffCostFunction(point[0],
                                                             point[1]),
          huber_loss, init_para);
    }

    ceres::Solver::Options options;
    options.max_num_iterations = 10;
    ceres::Solver::Summary summary;
    ceres::Solve(options, &problem, &summary);

    const double radius = init_para[2] * init_para[2];
    center_position[0] = init_para[0];
    center_position[1] = init_para[1];

    if (radius > options_.reflector_radius() - options_.radius_error() &&
        radius < options_.reflector_radius() + options_.radius_error()) {
      column_centers.emplace_back(centers.at(i));
      column_points_count.emplace_back(count.at(i));
    } else {
      LOG_EVERY_N(WARNING, 10) << "radius: " << radius;
    }
  }

  centers.clear();
  count.clear();
  centers = column_centers;
  count = column_points_count;
}

void ReflectorDetect::ExtractReflectorFeature(
    const sensor::PointCloud& reflector_raw_data,
    sensor::PointCloud* reflector_data_ptr) {
  sensor::PointCloud& reflector_data = *reflector_data_ptr;
  sensor::PointCloud reflector_centers;
  std::vector<int> center_counts;
  ClusterPoints(reflector_raw_data, &reflector_centers, &center_counts);

  reflector_data.clear();
  for (size_t i = 0u; i < center_counts.size(); ++i) {
    if (center_counts[i] >= options_.min_reflector_counter_per_cluster()) {
      reflector_data.push_back(reflector_centers[i]);
    }
  }
}

void ReflectorDetect::ReflectorMatching(
    const std::map<int, LandmarkPtr>& reflectors_map,
    const transform::Rigid2d& T_local_tracking,
    const sensor::PointCloud& reflector_feature_in_tracking,
    std::vector<int>* match_reflector_ids_ptr,
    std::vector<double>* match_min_dists_ptr) {
  if (reflector_feature_in_tracking.size() < 1) {
    return;
  }
  std::vector<int>& match_reflector_ids = *match_reflector_ids_ptr;
  std::vector<double>& match_min_dists = *match_min_dists_ptr;
  match_reflector_ids =
      std::vector<int>(reflector_feature_in_tracking.size(), -1);

  std::vector<Eigen::Vector2f> points_target;
  std::vector<Eigen::Vector2f> points_source;

  for (auto& point : reflector_feature_in_tracking) {
    auto reflector_in_local =
        T_local_tracking.cast<float>() * point.position.head<2>();
    points_source.push_back(reflector_in_local);
  }

  std::vector<int> landmark_ids;

  for (const auto& kv : reflectors_map) {
    const auto& reflector_id = kv.first;
    points_target.push_back(kv.second->p_local_landmark.cast<float>());
    landmark_ids.push_back(reflector_id);
  }

  std::vector<std::pair<int, int>> matches;

  for (size_t idx_source = 0u; idx_source < points_source.size();
       ++idx_source) {
    float min_dist = 100;
    int min_idx = -1;
    for (size_t idx_target = 0u; idx_target < points_target.size();
         ++idx_target) {
      const float dist =
          (points_target[idx_target] - points_source[idx_source]).norm();
      if (dist < min_dist) {
        min_dist = dist;
        min_idx = idx_target;
      }
    }
    if (min_idx != -1 && min_dist < options_.match_distance_threshold()) {
      matches.push_back(std::make_pair(min_idx, static_cast<int>(idx_source)));
    } else {
      LOG_EVERY_N(INFO, 10)
          << "No nearby reflector matched, min dist: " << min_dist;
    }
    match_min_dists.emplace_back(min_dist);
  }

  for (const auto& match : matches) {
    match_reflector_ids[match.second] = landmark_ids[match.first];
  }
}

void ReflectorDetect::ReflectorMatching(
    const std::map<std::string, Eigen::Vector3d>& reflectors_nearby,
    const transform::Rigid3d& T_global_tracking,
    const sensor::PointCloud& reflector_feature_in_tracking,
    std::vector<std::string>* match_reflector_ids_ptr,
    std::vector<double>* match_min_dists_ptr) {
  if (reflector_feature_in_tracking.size() < 1) {
    return;
  }
  std::vector<std::string>& match_reflector_ids = *match_reflector_ids_ptr;
  std::vector<double>& match_min_dists = *match_min_dists_ptr;
  match_reflector_ids =
      std::vector<std::string>(reflector_feature_in_tracking.size(), "");

  std::vector<Eigen::Vector2f> points_target;
  std::vector<Eigen::Vector2f> points_source;

  for (auto& point : reflector_feature_in_tracking) {
    auto reflector_in_global = T_global_tracking.cast<float>() * point;
    points_source.push_back(reflector_in_global.position.head<2>());
  }

  std::vector<std::string> landmark_ids;
  for (const auto& kv : reflectors_nearby) {
    const auto& reflector_id = kv.first;
    const Eigen::Vector3d& landmark_position = kv.second;
    points_target.push_back(landmark_position.cast<float>().head<2>());
    landmark_ids.push_back(reflector_id);
  }

  std::vector<std::pair<int, int>> matches;

  for (size_t idx_source = 0u; idx_source < points_source.size();
       ++idx_source) {
    float min_dist = 100;
    int min_idx = -1;
    for (size_t idx_target = 0u; idx_target < points_target.size();
         ++idx_target) {
      const float dist =
          (points_target[idx_target] - points_source[idx_source]).norm();
      if (dist < min_dist) {
        min_dist = dist;
        min_idx = idx_target;
      }
    }
    if (min_idx != -1 && min_dist < options_.match_distance_threshold()) {
      matches.push_back(std::make_pair(min_idx, static_cast<int>(idx_source)));
    } else {
      LOG_EVERY_N(INFO, 10)
          << "No nearby reflector matched, min dist: " << min_dist;
    }
    match_min_dists.emplace_back(min_dist);
  }

  for (const auto& match : matches) {
    match_reflector_ids[match.second] = landmark_ids[match.first];
  }
}

}  // namespace mapping
}  // namespace slam
